from .Matrix import *


def solve(a: Matrix, b: Vector):
    b = vec_from_mat(b, vertical=True)
    assert a.shape[1] == len(b)
    return vec_from_array(np.linalg.solve(a.store_, b.store_))


def cg(a: Matrix, b: Vector, residual_max=1e-10):
    """对实对称正定矩阵可以使用共轭梯度法"""
    b = vec_from_mat(b, vertical=True)
    assert a.shape[1] == len(b)
    n = len(b)

    assert (a.transpose == a).all()

    x = Vector(n, vertical=True)
    new_r = b - a * x
    new_p = new_r

    for iter_time in range(3 * n):
        old_r = new_r
        old_p = new_p

        # calculate new_r
        if a_squared_norm(old_r, a) < 0:
            raise ValueError("a should be positive definite")
        alpha = squared_norm(old_r) / a_squared_norm(old_r, a)

        x += alpha * old_p
        if alpha * old_p.norm < residual_max:
            break

        new_r = old_r - alpha * a * old_p
        beta = new_r.squared_norm / old_r.squared_norm
        new_p = new_r + beta * old_p

        if new_r.norm < 1e-15 or new_p.norm < 1e-15:
            break
    return vec_from_mat(x)


def gmres(a: Matrix, b: Matrix):
    pass
